package com.thinkingcoder.tcsharding.engine;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
import com.thinkingcoder.tcsharding.demo.model.User;
import com.thinkingcoder.tcsharding.property.ShardingProperties;
import com.thinkingcoder.tcsharding.strategy.HashShardingStrategy;
import com.thinkingcoder.tcsharding.strategy.ShardingStrategy;

/**
 * Standard Sharding Engine.
 */
public class StandardShardingEngine implements ShardingEngine {

    /**
     * key：数据库名字<br/>
     * value: 对应多个表的名字
     */
    private final MultiValueMap<String, String> actualDatabaseNames = new LinkedMultiValueMap<>();
    /**
     * key：表名字<br/>
     * value: 对应多个数据库的名字
     */
    private final MultiValueMap<String, String> actualTableNames = new LinkedMultiValueMap<>();

    private final Map<String, ShardingStrategy> databaseStrategies = new HashMap<>();
    private final Map<String, ShardingStrategy> tableStrategies = new HashMap<>();

    public StandardShardingEngine(ShardingProperties properties) {
        properties.getTables().forEach((table, tableProperties) -> {
            tableProperties.getActualDataNodes().forEach(actualDataNode -> {
                String[] split = actualDataNode.split("\\.");
                String databaseName = split[0], tableName = split[1];
                actualDatabaseNames.add(databaseName, tableName);
                actualTableNames.add(tableName, databaseName);
            });

            // 基于表的分库策略
            databaseStrategies.put(table, new HashShardingStrategy(tableProperties.getDatabaseStrategy()));
            // 基于表的分表策略
            tableStrategies.put(table, new HashShardingStrategy(tableProperties.getTableStrategy()));
        });
    }

    @Override
    public ShardingResult sharding(String sql, Object[] args) {

        SQLStatement sqlStatement = SQLUtils.parseSingleMysqlStatement(sql);
        if (sqlStatement instanceof SQLInsertStatement sqlInsertStatement) {
            String table = sqlInsertStatement.getTableName().getSimpleName();
            Map<String, Object> shardingColumnsMap = new HashMap<>();
            List<SQLExpr> columns = sqlInsertStatement.getColumns();
            for (int i = 0; i < columns.size(); i++) {
                SQLIdentifierExpr columnExpr = (SQLIdentifierExpr)columns.get(i);
                String columnName = columnExpr.getName();
                shardingColumnsMap.put(columnName, args[i]);
            }

            ShardingStrategy databaseStrategy = databaseStrategies.get(table);
            String targetDatabase =
                databaseStrategy.doSharding(actualDatabaseNames.get(table), table, shardingColumnsMap);

            ShardingStrategy tableStrategy = tableStrategies.get(table);
            String targetTable = tableStrategy.doSharding(actualTableNames.get(table), table, shardingColumnsMap);

            System.out.println(" ===>>> target db.table = " + targetDatabase + "." + targetTable);
        }

        Object parameterObject = args[0];
        System.out.println(" ===> getObject sql statement: " + sql);

        int id = 0;
        if (parameterObject instanceof User user) {
            id = user.getId();
        } else if (parameterObject instanceof Integer uid) {
            id = uid;
        }

        return new ShardingResult(id % 2 == 0 ? "ds0" : "ds1", sql);
    }

}
